//
// Created by mac on 1/24/22.
//
#include <iostream>
#include <vector>

using namespace std;


vector<int> findDiagonalOrder(vector<vector<int>>& mat) {
    //斜上 row -1 col +1
    //斜下 row +1 col -1
    //方向改变 ：考虑边界条件row>=m或者col>=n判断在前
    //       3. row>=m row = m-1, col+=2,换方向
    //       4. col>=n, row+=2, col = n-1,换方向
    //       1. row<0 row=0, col不变，换方向
    //       2. col<0 row不变 col=0, 换方向

    int m = mat.size();
    int n = mat[0].size();
    int row = 0;
    int col = 0;
    int dirs[2][2] = {{-1, 1}, {1,-1}};
    int dir = 0;

    vector<int> res(m*n, 0);
    for (int i = 0; i <m*n ; ++i) {
        res[i] = mat[row][col];

        row = row + dirs[dir][0];
        col = col + dirs[dir][1];

        if(row>=m){
            row = m-1;
            col = col+2;
            dir = 1-dir;
        }
        if(col>=n){
            row+=2;
            col = n-1;
            dir = 1-dir;
        }
        if (row < 0) {
            row = 0;
            dir = 1 - dir;
        }
        if (col < 0) {
            col = 0;
            dir = 1 - dir;
        }
    }
    return res;
}